package cn.fxc;

import org.apache.commons.math3.distribution.ChiSquaredDistribution;
import org.apache.commons.math3.distribution.PoissonDistribution;
import org.apache.commons.math3.stat.inference.ChiSquareTest;

import java.util.*;

public class PoissonChiSquareTest {
    int[] samples;
    private double pValue;
    private boolean rejected;
    private long[] observed;
    private double[] expected;
    private double chiSquare;

    public int[] getSamples() {
        return samples;
    }

    public void setSamples(int[] samples) {
        this.samples = samples;
    }

    public void chiSquareTest(){

        // 2. 计算频数分布
        Map<Integer, Long> frequencyMap = calculateFrequency(samples);

        // 3. 提取观测频数(处理缺失的类别)
//        int maxCategory = frequencyMap.keySet().stream().max(Integer::compare).orElse(0);

        // 4. 计算λ参数(泊松分布均值)
        double lambda = Arrays.stream(samples).average().orElse(0);
        int n = samples.length;
        // 5. 创建泊松分布实例
        PoissonDistribution poisson = new PoissonDistribution(lambda);
        // 5. 获取所有可能的x值(已排序)
        Integer[] xValues = frequencyMap.keySet().toArray(new Integer[0]);


        // 5. 前向合并分组
        int firstPartEnd = 0;
        double cumProb = 0;
        for (int i = 0; i < xValues.length; i++) {
            cumProb = poisson.cumulativeProbability(xValues[i]);
            if (n * cumProb > 2) {
                firstPartEnd = i;
                break;
            }
        }

        // 6. 后向合并分组
        int lastPartStart = xValues.length - 1;
        double tailProb = 0;
        for (int i = xValues.length - 1; i >= 0; i--) {
            tailProb = 1 - poisson.cumulativeProbability(xValues[i] - 1);
            if (n * tailProb > 2) {
                lastPartStart = i;
                break;
            }
        }

        // 7. 构建分组边界
        List<Integer> boundaries = new ArrayList<>();
        boundaries.add(xValues[0]); // 最小值

        // 前向合并区
        if (firstPartEnd > 0) {
            boundaries.add(xValues[firstPartEnd]);
        }
        // 中间独立分组
        int partStart = firstPartEnd;
        int partEnd = partStart +1;
        while (partStart < lastPartStart) {
            while (partEnd < lastPartStart) {
                cumProb = poisson.cumulativeProbability(xValues[partStart], xValues[partEnd]);
                if (n * cumProb > 2) {

                    boundaries.add(xValues[partEnd]);
                    break;
                }
                partEnd++;
            }

            partStart = partEnd;
            partEnd = partStart + 1;
        }

        // 后向合并区
        if (boundaries.get(boundaries.size()-1)<xValues[lastPartStart]) {
            boundaries.add(xValues[lastPartStart]);
        }
        if (boundaries.size()>=3){
            //如果后向合并区前面那个区累积分布概率*N小于2，说明多分了一个区，但是后向分布区不能去掉，只能把后向合并区前一个去掉(即size-2那一个)
            if (n * poisson.cumulativeProbability(boundaries.get(boundaries.size()-2),boundaries.get(boundaries.size()-1))<2){
                boundaries.remove(boundaries.size()-2);
            }
        }

        boundaries.add(Integer.MAX_VALUE); // 通过添加一个最大值，统一>=Xn与<Xn的运算

        // 8. 计算分组频数
        List<Long> observedList = new ArrayList<>();
        List<Double> expectedList = new ArrayList<>();

        for (int i = 0; i < boundaries.size() - 1; i++) {
            int lower = boundaries.get(i);
            int upper = boundaries.get(i + 1);

            // 计算观测频数
            long obs = 0;
            for (Map.Entry<Integer, Long> entry : frequencyMap.entrySet()) {
                if (entry.getKey() >= lower && entry.getKey() < upper) {
                    obs += entry.getValue();
                }
            }
            observedList.add(obs);


            // 计算期望频数
            double exp = n * ((poisson.cumulativeProbability(upper-1) -
                    (lower > 0 ? poisson.cumulativeProbability(lower - 1) : 0)));
            expectedList.add(exp);
        }

        // 转换为数组
        observed = observedList.stream().mapToLong(l -> l).toArray();
        expected = expectedList.stream().mapToDouble(d -> d).toArray();

        // 8. 执行卡方检验
        ChiSquareTest chiSquareTest = new ChiSquareTest();
        chiSquare = chiSquareTest.chiSquare(expected, observed);

        ChiSquaredDistribution chiSquaredDistribution = new ChiSquaredDistribution(observed.length-1-1);

        pValue = 1-chiSquaredDistribution.cumulativeProbability(chiSquare);
        rejected = pValue <0.05;

        // 9. 输出结果
        System.out.println("原始样本值: " + Arrays.toString(samples));
        System.out.println("频数分布: " + frequencyMap);
        System.out.println("λ (lambda) 估计值: " + lambda);

        System.out.println("分组边界: " + Arrays.toString(boundaries.toArray()));
        System.out.println("合并后的观测频数: " + Arrays.toString(observed));
        System.out.println("合并后的期望频数: " + Arrays.toString(expected));
        System.out.println("卡方统计量: " + chiSquare);
        System.out.println("P值: " + pValue);
        System.out.println("在0.05显著性水平下是否拒绝原假设? " + (rejected ? "拒绝" : "不拒绝"));
        System.out.println("自由度: " + (observed.length - 1 - 1)); // 类别数 - 1 - 估计的参数数
    }

    // 计算频数分布
    private static Map<Integer, Long> calculateFrequency(int[] samples) {
        // 2. 计算频数分布(使用TreeMap自动排序)
        Map<Integer, Long> freqMap = new TreeMap<>();
        for (int sample : samples) {
            freqMap.put(sample, freqMap.getOrDefault(sample, 0L) + 1);
        }

        return freqMap;
    }


    public static void main(String[] args) {

        int[] samples = {2,2,3,3,2,4,4,2,1,1,1,4,4,3,0,4,3,2,3,3,4,1,3,1,4,3,2,2,1,2,0,2,3,2,3};
        PoissonChiSquareTest poissonChiSquareTest = new PoissonChiSquareTest();
        poissonChiSquareTest.setSamples(samples);
        poissonChiSquareTest.chiSquareTest();
    }

}
